package wt.ai.chatagentservice.controller;


import com.alibaba.fastjson.JSONObject;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.service.TokenStream;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;
import wt.ai.chatagentservice.domain.ChatAgentMessage;
import wt.ai.chatagentservice.domain.ChatDataMessage;
import wt.ai.chatagentservice.domain.R;
import wt.ai.chatagentservice.service.Assistant;
import wt.ai.chatagentservice.service.ToolService;

import java.io.IOException;

@RestController
public class ChatAgentController {


    @Autowired
    private Assistant assistant;
    @Autowired
    private ToolService toolService;

    @GetMapping("/chat")
    public R model(@RequestParam(value = "message", defaultValue = "Hello") String message) {

        return R.ok(assistant.chat(message));
    }


    @PostMapping("/chatStream")
    public Flux<String> chatStream(@RequestBody ChatAgentMessage chatAgentMessage) {

        Flux<String> stringFlux = assistant.chatStream(chatAgentMessage.getMessage());

        return stringFlux;
    }



    @PostMapping("/chatWithThink")
    public SseEmitter chatWithThink(@RequestBody ChatAgentMessage chatAgentMessage) {
        SseEmitter emitter = new SseEmitter(0L);


        TokenStream tokenStream = assistant.chatWithThink(chatAgentMessage.getMessage());

        tokenStream.onPartialThinking(partialThinking -> {
                   try {
                       ChatDataMessage chatDataMessage = new ChatDataMessage();
                       chatDataMessage.setContent(partialThinking.text());
                       chatDataMessage.setType("thinking");
                       chatDataMessage.setDone(false);
                       emitter.send(SseEmitter.event()
                               .name("thinking")
                               .data( JSONObject.toJSONString(chatDataMessage)));
                   } catch (IOException e) {
                       throw new RuntimeException(e);
                   }
               }).onToolExecuted(tool -> {
                    System.out.println(tool);
                })
                .onPartialResponse(partial -> {
                    try {
                        ChatDataMessage chatDataMessage = new ChatDataMessage();
                        chatDataMessage.setContent(partial);
                        chatDataMessage.setType("partialResponse");
                        emitter.send(SseEmitter.event()
                                .name("partialResponse")
                                .data( JSONObject.toJSONString(chatDataMessage)));
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                })
                 .onRetrieved(contents -> {

                })
                .onCompleteResponse(response -> {

                    try {
//                        ChatDataMessage chatDataMessage = new ChatDataMessage();
//                        chatDataMessage.setContent(response.aiMessage().thinking());
//                        chatDataMessage.setType("thinking");
//                        chatDataMessage.setDone(true);
//                        emitter.send(SseEmitter.event()
//                                .name("thinking")
//                                .data( JSONObject.toJSONString(chatDataMessage)));
                    } catch (Exception e) {
                        throw new RuntimeException(e);
                    }



                    emitter.complete();
                })
                .onError(error -> {
                    emitter.completeWithError(error);
                })
                .start();


        return emitter;
    }


}
